from pathlib import Path

from sklearn.neighbors import KNeighborsClassifier
from sklearn.preprocessing import LabelEncoder
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report

from SimpleProcessor import Simple_Processor
from SimpleDataSetLoader import Simple_DataSet_Loader


def knn(images_path: Path, neighbor: int, jobs: int = -1):
    # images_path = Path('../dataset/dog_cat')
    sp = Simple_Processor(32, 32)
    sdl = Simple_DataSet_Loader([sp])
    data, labels = sdl.load([str(i) for i in images_path.iterdir()])
    data = data.reshape((data.shape[0], 3072))
    print(f'[info]: 特征矩阵--{(data.nbytes / (1024 * 1000.0)):.2f}Mb')
    le = LabelEncoder()
    labels = le.fit_transform(labels)
    train_x, test_x, train_y, test_y = train_test_split(data, labels,
                                                        test_size=0.25,
                                                        random_state=42)
    print('[info]: 评估knn模型...')
    knn_model = KNeighborsClassifier(n_neighbors=neighbor,
                                     n_jobs=jobs)  # n_jobs用于cpu占比，-1为所有核心
    knn_model.fit(train_x, train_y)
    print(classification_report(test_y, knn_model.predict(test_x),
                                target_names=le.classes_))


if __name__ == '__main__':
    knn(Path('../../dataset/dog_cat'), 2)
